from scipy.io import loadmat as load
import numpy as np
import matplotlib.pyplot as plt

num_labs = 10
image_size = 32


def reformat(samples, labels):
    samples = np.transpose(samples, (3, 0, 1, 2))
    print(labels.shape)
    labels = np.array([x[0] for x in labels])
    one_hot_labels = []
    print(len(labels))
    for num in labels:
        one_hot = [0.0] * 10
        if num == 10:
            one_hot[0] = 1.0
        else:
            one_hot[num] = 1.0
        one_hot_labels.append(one_hot)
    labels = np.array(one_hot_labels).astype(np.float32)
    return samples, labels


def normalize(samples):
    a = np.add.reduce(samples, keepdims=True, axis=3)
    a = a / 3.0
    return a/128 - 1


def destribution(labels, name):
    pass


def inspect(dataset, labels, i):
    print(labels[i])
    plt.imshow(dataset[i])
    plt.show()
    pass

train_data = load('./train_32x32.mat')
test_data = load('./test_32x32.mat')
# extra_data = load('./extra_32x32.mat')

train_samples = train_data['X']
train_labels = train_data['y']
test_samples = train_data['X']
test_labels = train_data['y']
_train_samples, _train_labels = reformat(train_samples, train_labels)
_test_samples, _test_labels = reformat(test_samples, test_labels)

if __name__ == "__main__":
    _train_samples = normalize(_train_samples)
    inspect(_train_samples, _train_labels, 8823)
    pass
